deepmd/jax/jax2tf/region.py (3)

29-48: Extract common shape validation logic.

Similar to phys2inter, this function would benefit from shape validation. Consider extracting the validation into a shared utility function to avoid code duplication.

+def _validate_coord_cell_shapes(coord: tnp.ndarray, cell: tnp.ndarray, name: str):
+    coord_shape = tf.shape(coord)
+    cell_shape = tf.shape(cell)
+    tf.debugging.assert_equal(
+        coord_shape[-1],
+        cell_shape[-1],
+        message=f"{name}: Last dimension of coord and cell must match"
+    )
+    tf.debugging.assert_equal(
+        cell_shape[-2:],
+        tf.constant([3, 3]),
+        message=f"{name}: Cell must be a 3x3 matrix"
+    )

 def inter2phys(...):
+    _validate_coord_cell_shapes(coord, cell, "inter2phys")
     return tnp.matmul(coord, cell)

---

51-72: Optimize matrix operations for better performance.

The current implementation performs two matrix multiplications by calling phys2inter and inter2phys. This could be optimized by caching the inverse cell matrix and reusing it.

 def normalize_coord(
     coord: tnp.ndarray,
     cell: tnp.ndarray,
 ) -> tnp.ndarray:
+    rec_cell = tf.linalg.inv(cell)
-    icoord = phys2inter(coord, cell)
+    icoord = tnp.matmul(coord, rec_cell)
     icoord = tnp.remainder(icoord, 1.0)
-    return inter2phys(icoord, cell)
+    return tnp.matmul(icoord, cell)

---

75-93: Add validation for cell tensor in to_face_distance.

The function reshapes the input tensor but doesn't validate that the final dimensions are correct.

 def to_face_distance(
     cell: tnp.ndarray,
 ) -> tnp.ndarray:
     cshape = tf.shape(cell)
+    tf.debugging.assert_equal(
+        cshape[-2:],
+        tf.constant([3, 3]),
+        message="Cell must be a 3x3 matrix"
+    )
     dist = b_to_face_distance(tnp.reshape(cell, [-1, 3, 3]))
     return tnp.reshape(dist, tf.concat([cshape[:-2], [3]], axis=0))

deepmd/jax/utils/serialization.py (2)

72-76: Simplify ghost atoms logic using ternary operator

The logic is correct, but it can be more concise using a ternary operator.

-            if has_ghost_atoms:
-                nghost_ = nghost
-            else:
-                nghost_ = 0
+            nghost_ = nghost if has_ghost_atoms else 0

🧰 Tools

🪛 Ruff

72-75: Use ternary operator nghost_ = nghost if has_ghost_atoms else 0 instead of if-else-block

Replace if-else-block with nghost_ = nghost if has_ghost_atoms else 0

(SIM108)

---

Tests needed for new serialization configurations

The implementation introduces new serialization configurations for ghost atoms and atomic virial combinations, but no corresponding test coverage was found. Tests should be added to verify:

• Serialization with ghost atoms enabled/disabled
• Serialization with atomic virial calculation enabled/disabled
• Combined scenarios of both configurations

Specific areas needing test coverage:

• exported_whether_do_atomic_virial() function behavior
• Serialization of all four combinations stored in @variables
• Proper handling of ghost atoms in the exported models

🔗 Analysis chain

Line range hint 56-119: Verify test coverage for all configurations

While the implementation looks solid, we should ensure proper test coverage for all combinations of:

• Atomic virial calculation (True/False)
• Ghost atoms presence (True/False)

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for test files covering these scenarios
rg -l "exported_whether_do_atomic_virial" --type python | grep "test"

Length of output: 98

---

Script:

#!/bin/bash
# Search for test files covering these scenarios without file type filter
rg -l "exported_whether_do_atomic_virial" | grep "test"

Length of output: 53

---

Script:

#!/bin/bash
# Search for test files that might cover the serialization functionality
rg -l "stablehlo.*ghost|stablehlo.*atomic_virial" | grep "test"

Length of output: 63

---

Script:

#!/bin/bash
# Look for test files in the deepmd directory
fd "test.*\.py$" deepmd/

Length of output: 83

---

Script:

#!/bin/bash
# Search for test files in the jax directory specifically
fd "test.*\.py$" deepmd/jax/

# Look for any serialization tests
rg -l "serialize|serializ" deepmd/jax/ | grep "test"

# Check content of test files for related functionality
rg "atomic_virial|ghost" deepmd/jax/ -g "test*.py" -C 2

Length of output: 141

source/tests/consistent/io/test_io.py (2)

186-202: Add documentation for nopbc test cases.

While the implementation is correct, it would be helpful to add a comment explaining that passing None as the box parameter indicates no periodic boundary conditions. This would make the test's intention more explicit.

+ # Test evaluation without periodic boundary conditions by passing None as box
 ret = deep_eval.eval(
     self.coords,
     None,
     self.atype,
     fparam=fparam,
     aparam=aparam,
 )

---

The review comment's suggestion for additional test coverage is valid.

While the codebase has logic to handle zero ghost atoms in the JAX serialization code, there are no explicit test cases verifying this scenario. The existing tests in test_virtual_type.py only cover non-zero ghost atoms (nghost=10). Adding test cases for zero ghost atoms would help ensure the polymorphic_shape fix works correctly for this edge case.

• Add test cases in source/tests/consistent/io/test_io.py to verify model behavior with zero ghost atoms
• Consider adding similar test coverage in source/tests/tf/test_virtual_type.py

🔗 Analysis chain

Line range hint 145-218: Consider adding test cases for zero ghost atoms.

While the test coverage for nopbc is comprehensive, consider adding specific test cases to verify the fix for zero ghost atoms mentioned in the PR objectives. This would ensure that the polymorphic_shape issue is properly addressed.

---

Based on the search results, I can see that there are existing tests for ghost atoms in source/tests/tf/test_virtual_type.py and source/tests/tf/common.py that cover cases with non-zero ghost atoms (e.g., nghost = 10 and nghost = 4). The codebase also has specific handling for zero ghost atoms in JAX serialization (deepmd/jax/jax2tf/serialization.py and deepmd/jax/utils/serialization.py).

Let me verify if there are specific test cases for zero ghost atoms.

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Description: Check if there are any existing tests for zero ghost atoms
# Test: Search for test cases involving ghost atoms
rg -i "ghost.*atom|nghost" --type py

Length of output: 8281

---

Script:

#!/bin/bash
# Check for test cases that explicitly test zero ghost atoms
rg -B2 -A2 "nghost.*=.*0|has_ghost_atoms.*=.*false" --type py

Length of output: 1487

deepmd/jax/jax2tf/tfmodel.py (1)

144-147: Consider using a ternary operator for better readability

The call selection can be more concise.

-        if do_atomic_virial:
-            call = self._call_atomic_virial
-        else:
-            call = self._call
+        call = self._call_atomic_virial if do_atomic_virial else self._call

🧰 Tools

🪛 Ruff

144-147: Use ternary operator call = self._call_atomic_virial if do_atomic_virial else self._call instead of if-else-block

Replace if-else-block with call = self._call_atomic_virial if do_atomic_virial else self._call

(SIM108)

deepmd/jax/jax2tf/make_model.py (4)

50-72: Document all function parameters in the docstring

The parameters call_lower, rcut, sel, mixed_types, and model_output_def are not documented in the docstring. Please add descriptions for these parameters to enhance clarity and maintainability.

Example addition to the docstring:

     """
     Return model prediction from lower interface.

     Parameters
     ----------
+    call_lower
+        A callable that accepts extended coordinates and other parameters, returning model predictions as a dictionary.
+    rcut
+        Cut-off radius for neighbor list construction.
+    sel
+        List of integers specifying selected atom types.
+    mixed_types
+        Boolean indicating whether to treat atom types as mixed.
+    model_output_def
+        Definition of the model outputs.
     coord
         The coordinates of the atoms.
         shape: nf x (nloc x 3)

---

75-76: Improve variable naming for better readability

Using aliases like cc, bb, fp, ap for coord, box, fparam, aparam reduces code readability. It's clearer to use the original variable names throughout the function.

Apply this diff to enhance readability:

- cc, bb, fp, ap = coord, box, fparam, aparam
- del coord, box, fparam, aparam
+ # Use original variable names throughout the function

---

75-76: Remove unnecessary variable deletion

The del statements for coord, box, fparam, and aparam might be unnecessary. Python's garbage collector handles memory management, and removing these lines can simplify the code.

Apply this diff to remove unnecessary code:

  cc, bb, fp, ap = coord, box, fparam, aparam
- del coord, box, fparam, aparam

---

93-93: Clarify the logic of distinguish_types parameter

The expression distinguish_types=not mixed_types may be confusing. Consider renaming mixed_types to ignore_atom_types or adding a comment to clarify that distinguish_types is the inverse of mixed_types.

deepmd/jax/jax2tf/transform_output.py (3)

54-54: Remove unused variable mldims.

The variable mldims is assigned but never used, which could lead to confusion. Please remove it to keep the code clean.

🧰 Tools

🪛 Ruff

54-54: Local variable mldims is assigned to but never used

Remove assignment to unused variable mldims

(F841)

🪛 GitHub Check: CodeQL

[notice] 54-54: Unused local variable
Variable mldims is not used.

---

83-83: Replace assert with explicit exception handling.

Using assert statements for control flow may not be ideal in production code because they can be disabled with optimization flags. Consider raising a specific exception with a clear error message to handle cases where vdef.r_differentiable is False while vdef.c_differentiable is True.

---

41-44: Enhance the docstring for communicate_extended_output.

The current docstring is brief. Providing detailed descriptions of the parameters, return values, and any important computational steps would improve maintainability and clarity for other developers.

deepmd/jax/jax2tf/nlist.py (6)

2-4: Simplify import statement by removing unnecessary parentheses

The parentheses around Union are unnecessary when importing a single item.

Apply this diff to simplify the import:

-from typing import (
-    Union,
-)
+from typing import Union

---

9-11: Simplify import statement by removing unnecessary parentheses

The parentheses around to_face_distance are unnecessary when importing a single item.

Apply this diff to simplify the import:

-from .region import (
-    to_face_distance,
-)
+from .region import to_face_distance

---

14-14: Remove reference to 'chatgpt' in comment for professionalism

It's better to avoid mentioning specific tools like 'chatgpt' in code comments for professionalism and maintainability.

Apply this diff to update the comment:

-## translated from torch implementation by chatgpt
+## Translated from torch implementation

---

28-28: Correct typo in docstring: 'exptended' should be 'extended'

There's a typo in the parameter description for coord.

Apply this diff to fix the typo:

-        exptended coordinates of shape [batch_size, nall x 3]
+        extended coordinates of shape [batch_size, nall x 3]

---

154-154: Correct typo in docstring: 'peridoc' should be 'periodic'

There's a typo in the function description.

Apply this diff to fix the typo:

-    """Extend the coordinates of the atoms by appending peridoc images.
+    """Extend the coordinates of the atoms by appending periodic images.

---

147-147: Remove reference to 'chatgpt' in comment for professionalism

As before, it's better to avoid mentioning specific tools like 'chatgpt' in code comments.

Apply this diff to update the comment:

-## translated from torch implementation by chatgpt
+## Translated from torch implementation

deepmd/jax/model/hlo.py (2)

183-192: Consider adding comments to clarify the conditional logic in call_lower.

The conditional logic in call_lower now depends on the dimensions of extended_coord and nlist, as well as the do_atomic_virial flag. Adding comments to explain the reasoning behind these conditions can improve code readability and maintainability.

For example:

 def call_lower(
     self,
     extended_coord: jnp.ndarray,
     extended_atype: jnp.ndarray,
     nlist: jnp.ndarray,
     mapping: Optional[jnp.ndarray] = None,
     fparam: Optional[jnp.ndarray] = None,
     aparam: Optional[jnp.ndarray] = None,
     do_atomic_virial: bool = False,
 ):
+    # Determine if ghost atoms are present based on the shape of extended_coord and nlist
     if extended_coord.shape[1] > nlist.shape[1]:
+        # Case with ghost atoms
         if do_atomic_virial:
             call_lower = self._call_lower_atomic_virial
         else:
             call_lower = self._call_lower
     else:
+        # Case without ghost atoms
         if do_atomic_virial:
             call_lower = self._call_lower_atomic_virial_no_ghost
         else:
             call_lower = self._call_lower_no_ghost
     return call_lower(
         extended_coord,
         extended_atype,

---

183-192: Ensure unit tests cover all branches of the new conditional logic.

With the updated call_lower method introducing additional branches, it's important to verify that all scenarios are properly tested. This includes cases:

• With and without ghost atoms (i.e., when extended_coord.shape[1] > nlist.shape[1] and when it's not).
• With do_atomic_virial set to both True and False.

deepmd/jax/jax2tf/serialization.py (1)

55-58: Simplify assignment with a ternary operator

Consider replacing the if-else block with a ternary operator for conciseness and readability.

Apply this diff to implement the suggestion:

             # nghost >= 1 is assumed if there is
             # other workaround does not work, such as
             # nall; nloc + nghost - 1
-            if has_ghost_atoms:
-                nghost = "nghost"
-            else:
-                nghost = "0"
+            nghost = "nghost" if has_ghost_atoms else "0"

🧰 Tools

🪛 Ruff

55-58: Use ternary operator nghost = "nghost" if has_ghost_atoms else "0" instead of if-else-block

Replace if-else-block with nghost = "nghost" if has_ghost_atoms else "0"

(SIM108)

Configuration used: CodeRabbit UI
Review profile: CHILL

#!/bin/bash
# Check TensorFlow documentation or issues related to experimental.numpy stability
gh search issues "tensorflow experimental.numpy stability" --repo tensorflow/tensorflow --json title,body --jq '.[] | select(.title | contains("experimental.numpy") or contains("tnp")) | {title: .title, body: .body}'

#!/bin/bash
# Let's check the actual usage of tnp in the codebase to understand its impact
rg "import tensorflow.experimental.numpy as tnp" -A 3

# Also check for any direct usage of tnp
rg "\btnp\." -B 1 -A 1

# Look for any configuration or setup related to numpy behavior
rg "enable_numpy_behavior|set_numpy_enabled"

#!/bin/bash
# Description: Check for the presence of new variables in HLO files
# Expected: Files containing these new variables

# Search for files containing the new variable names
rg -l "stablehlo_no_ghost|stablehlo_atomic_virial_no_ghost" --type py

# Search for test files that verify these variables
rg -l "test.*stablehlo.*ghost" --type py

#!/bin/bash
# Let's examine the implementation and usage of these variables in the found files

# Check hlo.py implementation
echo "=== Content from deepmd/jax/model/hlo.py ==="
rg -A 5 -B 5 "stablehlo_no_ghost|stablehlo_atomic_virial_no_ghost" deepmd/jax/model/hlo.py

echo -e "\n=== Content from deepmd/jax/utils/serialization.py ==="
rg -A 5 -B 5 "stablehlo_no_ghost|stablehlo_atomic_virial_no_ghost" deepmd/jax/utils/serialization.py

# Check if there are any tests that might not directly mention "stablehlo" and "ghost"
echo -e "\n=== Related test files ==="
fd "test.*\.py" | xargs rg -l "no_ghost"

#!/bin/bash
# Description: Verify the signature of `call_lower`.

# Expected: `call_lower` accepts the specified arguments and matches the Callable type hint.

# Since `call_lower` is passed as a parameter, search for typical usages or definitions
rg -A 5 'def call_lower' | rg 'def call_lower\('

# Alternatively, check for any type hints or comments indicating the expected parameters
rg -A 5 'call_lower' | rg 'Callable\[['

#!/bin/bash
# Description: Check how `do_atomic_virial` is used in `communicate_extended_output`.

# Expected: The flag is used to conditionally process atomic virial calculations.

# Find the definition of `communicate_extended_output` and its usage of `do_atomic_virial`
rg -A 15 'def communicate_extended_output' | rg 'do_atomic_virial'

#!/bin/bash
# Description: Check if `extend_coord_with_ghosts` handles `box` being `None`.

# Expected: No errors or exceptions when `box` is `None`.

# Find the definition of `extend_coord_with_ghosts`
rg -A 10 'def extend_coord_with_ghosts' | grep 'box'

# Search for any checks or handling of `box` being `None` within the function
rg -A 20 'def extend_coord_with_ghosts' | rg 'if box is None|box is not None'

#!/bin/bash
# Let's try a different approach to find the implementation of extend_coord_with_ghosts

# Search for the function definition with more context
rg -A 20 "def extend_coord_with_ghosts"

# Also search in case it's defined with different spacing
rg -A 20 "def\s+extend_coord_with_ghosts"

# Try finding any usage of this function to understand the context
rg "extend_coord_with_ghosts\("

#!/bin/bash
# Let's check the implementation in jax2tf/nlist.py since that's the one being used in make_model.py
rg -A 50 "def extend_coord_with_ghosts" deepmd/jax/jax2tf/nlist.py

# Also check if there's any handling of None cell/box in the dpmodel implementation since it has Optional typing
rg -A 50 "def extend_coord_with_ghosts" deepmd/dpmodel/utils/nlist.py

source/tests/jax/jax2tf/__init__.py (1)

1-8: Add missing TensorFlow configurations

Based on the AI summary, this file should include TensorFlow's experimental numpy configuration and eager execution check.

Consider adding these configurations:

 # SPDX-License-Identifier: LGPL-3.0-or-later
 import pytest
+import tensorflow as tf
+import tensorflow.experimental.numpy as tnp
 
 from ...utils import (
     DP_TEST_TF2_ONLY,
 )
 
+if not tf.executing_eagerly():
+    raise RuntimeError("This module must be run in eager execution mode")
+
+tnp.experimental_enable_numpy_behavior()
+
 pytestmark = pytest.mark.skipif(not DP_TEST_TF2_ONLY, reason="TF2 conflicts with TF1")

source/tests/jax/jax2tf/test_region.py (3)

4-4: Consider the stability implications of using experimental numpy.

While using tensorflow.experimental.numpy is acceptable in tests, be aware that its API might change in future TensorFlow versions. Consider adding a comment explaining why the experimental version is preferred over the stable numpy.

---

18-22: Add documentation for cell initialization.

The cell initialization involves complex shape transformations. Consider adding comments to explain:

1. Why the cell is initialized with this specific 3D structure
2. The purpose of reshaping to [1, 1, -1, 3]
3. The significance of tiling to [4, 5, 1, 1]

---

25-34: Enhance test coverage with edge cases.

While the current test with random inputs is good, consider adding specific test cases for:

1. Zero vectors
2. Unit vectors
3. Extreme values
4. Negative coordinates

This would ensure the transformation handles all possible scenarios correctly.

source/tests/jax/jax2tf/test_nlist.py (2)

17-44: Add docstrings to improve test documentation.

The test class and setUp method would benefit from docstrings explaining:

• The purpose of the test suite
• The structure of the test data
• The meaning of key parameters (nf, nloc, ns, etc.)

Example addition:

 class TestNeighList(tf.test.TestCase):
+    """Test suite for neighbor list construction and coordinate extension.
+    
+    Tests the functionality of build_neighbor_list and extend_coord_with_ghosts
+    with realistic molecular simulation data.
+    """
     def setUp(self):
+        """Initialize test data including coordinates, cell, and reference lists.
+        
+        Sets up:
+        - nf (int): Number of frames (3)
+        - nloc (int): Number of local atoms (3)
+        - ns (int): Number of shifts (75)
+        - cell (array): Simulation cell parameters
+        - coordinates and atom types
+        """

---

45-65: Extract magic numbers into named constants.

The test uses magic numbers for array comparisons. Consider extracting these into named constants at the class level for better maintainability.

 class TestNeighList(tf.test.TestCase):
+    # Constants for array comparisons
+    INVALID_NEIGHBOR = -1
+
     def test_build_notype(self):
-        nlist_mask = nlist[0] == -1
+        nlist_mask = nlist[0] == self.INVALID_NEIGHBOR

Configuration used: CodeRabbit UI
Review profile: CHILL

#!/bin/bash
# Search for connections between region transformations and ghost atoms
rg -l "ghost.*region|region.*ghost" --type py

# Look for usage of these transformation functions in neighbor list construction
ast-grep --pattern 'inter2phys|to_face_distance' 

#!/bin/bash
# Search for files containing ghost atoms implementation
rg -l "ghost.*atom" --type py

# Search for implementations using region transformations
rg -A 5 "inter2phys|to_face_distance" --type py

# Look for neighbor list implementations
rg -l "neighbor.*list|neigh.*list" --type py

# Check the actual implementation files for region.py
fd region.py

#!/bin/bash
# Description: Verify that the test files exist and contain relevant test cases

# Test 1: Check if the test files exist
echo "Checking test files existence..."
fd -t f "test_io.py$" "source/tests/consistent/io"
fd -t f "test_nlist.py$" "source/tests/jax/jax2tf"

# Test 2: Check test content for neighbor list and nopbc related tests
echo "Checking test content..."
rg -l "test.*nlist|test.*pbc" "source/tests/jax/jax2tf"

#!/bin/bash
# Let's check the actual content of these test files to verify coverage

echo "Content of test_io.py:"
cat source/tests/consistent/io/test_io.py

echo -e "\nContent of test_nlist.py:"
cat source/tests/jax/jax2tf/test_nlist.py

# Also check for any other test files in jax2tf directory
echo -e "\nOther test files in jax2tf directory:"
fd -t f "^test.*\.py$" "source/tests/jax/jax2tf"

#!/bin/bash
# Check if similar matrix transformations are documented elsewhere
rg "tf.linalg.inv.*cell" -B 2 -A 2

source/tests/jax/jax2tf/test_region.py (2)

26-34: Add docstrings to explain the test class and cell structure.

The test setup is well-implemented, but would benefit from documentation explaining:

• The purpose of the TestRegion class
• The structure and significance of the cell array
• Why these specific dimensions (4, 5, -1, 3) are used

 class TestRegion(tf.test.TestCase):
+    """Tests for region transformation functions.
+    
+    The test class validates coordinate transformations and distance calculations
+    using a 4x5 batch of 3x3 cell matrices. Each cell represents a transformation
+    from internal to physical coordinates.
+    """
     def setUp(self):
+        """Sets up test fixtures with a batched cell array for coordinate transforms."""

---

35-44: Consider adding edge cases to the coordinate transformation test.

While the basic test is solid, consider adding test cases for:

• Zero coordinates
• Large coordinate values
• Negative coordinates
• Boundary values

 def test_inter_to_phys(self):
+    """Test internal to physical coordinate transformation with various cases."""
     rng = tf.random.Generator.from_seed(GLOBAL_SEED)
     inter = rng.normal(shape=[4, 5, 3, 3])
+    # Test regular case
     phys = inter2phys(inter, self.cell)
     for ii in range(4):
         for jj in range(5):
             expected_phys = tnp.matmul(inter[ii, jj], self.cell[ii, jj])
             self.assertAllClose(
                 phys[ii, jj], expected_phys, rtol=self.prec, atol=self.prec
             )
+    
+    # Test edge cases
+    edge_cases = [
+        tnp.zeros([4, 5, 3, 3]),  # Zero coordinates
+        tnp.ones([4, 5, 3, 3]) * 1e6,  # Large values
+        tnp.ones([4, 5, 3, 3]) * -1.0,  # Negative values
+    ]
+    for case in edge_cases:
+        phys = inter2phys(case, self.cell)
+        for ii in range(4):
+            for jj in range(5):
+                expected_phys = tnp.matmul(case[ii, jj], self.cell[ii, jj])
+                self.assertAllClose(
+                    phys[ii, jj], expected_phys, rtol=self.prec, atol=self.prec
+                )

source/tests/jax/jax2tf/test_nlist.py (2)

28-55: Add docstrings to explain test setup and data structure.

The test setup is comprehensive but would benefit from documentation explaining:

• The purpose and structure of the test data
• The meaning of magic numbers (e.g., nloc=3, ns=553)
• The format and meaning of ref_nlist array

Example docstring:

class TestNeighList(tf.test.TestCase):
    """Tests for neighbor list construction and coordinate extension.
    
    Test data structure:
    - nf: Number of frames
    - nloc: Number of local atoms per frame
    - ns: Number of shifts (5x5x3 grid)
    - ref_nlist: Reference neighbor list with format [...]
    """

---

77-98: Optimize test_build_type by reducing redundant operations.

The test method could be optimized by:

1. Caching the mapping operation results
2. Using vectorized operations instead of the loop

 def test_build_type(self):
     ecoord, eatype, mapping = extend_coord_with_ghosts(
         self.coord, self.atype, self.cell, self.rcut
     )
     nlist = build_neighbor_list(
         ecoord,
         eatype,
         self.nloc,
         self.rcut,
         self.nsel,
         distinguish_types=True,
     )
     self.assertAllClose(nlist[0], nlist[1])
     nlist_mask = nlist[0] == -1
     nlist_loc = mapping[0][nlist[0]]
     nlist_loc = tnp.where(nlist_mask, tnp.full_like(nlist_loc, -1), nlist_loc)
-    for ii in range(2):
-        self.assertAllClose(
-            tnp.sort(tnp.split(nlist_loc, self.nsel, axis=-1)[ii], axis=-1),
-            tnp.sort(tnp.split(self.ref_nlist, self.nsel, axis=-1)[ii], axis=-1),
-        )
+    # Split and sort once, then compare
+    nlist_split = tnp.sort(tnp.split(nlist_loc, self.nsel, axis=-1), axis=-1)
+    ref_split = tnp.sort(tnp.split(self.ref_nlist, self.nsel, axis=-1), axis=-1)
+    self.assertAllClose(nlist_split, ref_split)

Configuration used: CodeRabbit UI
Review profile: CHILL

 def test_to_face_dist(self):
+    """Test face distance calculations with validation."""
     cell0 = self.cell[0][0]
     vol = tf.linalg.det(cell0)
+    # Ensure cell is not singular
+    self.assertGreater(tf.abs(vol), 1e-10)
+
     # area of surfaces xy, xz, yz
     sxy = tf.linalg.norm(tnp.cross(cell0[0], cell0[1]))
     sxz = tf.linalg.norm(tnp.cross(cell0[0], cell0[2]))
     syz = tf.linalg.norm(tnp.cross(cell0[1], cell0[2]))
+    # Validate surface areas
+    self.assertGreater(sxy, 1e-10)
+    self.assertGreater(sxz, 1e-10)
+    self.assertGreater(syz, 1e-10)
+
     # vol / area gives distance
     dz = vol / sxy
     dy = vol / sxz
     dx = vol / syz
     expected = tnp.array([dx, dy, dz])
     dists = to_face_distance(self.cell)
-    for ii in range(4):
-        for jj in range(5):
-            self.assertAllClose(
-                dists[ii][jj], expected, rtol=self.prec, atol=self.prec
-            )
+    # Vectorized comparison
+    self.assertAllClose(
+        dists, tnp.broadcast_to(expected, [4, 5, 3]),
+        rtol=self.prec, atol=self.prec
+    )

source/tests/consistent/io/test_io.py (1)

149-149: Add docstring explaining nopbc test cases.

While the implementation of no periodic boundary condition tests is good, consider adding documentation to explain:

1. The purpose of these test cases
2. The expected behavior when box is None
3. The relationship with the PR's nopbc fixes

Example docstring addition:

def test_deep_eval(self):
+    """Test deep_eval with both periodic and non-periodic boundary conditions.
+    
+    This test verifies:
+    1. Regular evaluation with periodic boundary conditions
+    2. Evaluation without periodic boundaries (box=None)
+    3. Atomic-level outputs for both cases
+    """

Also applies to: 190-206

Configuration used: CodeRabbit UI
Review profile: CHILL

#!/bin/bash
# Description: Check if there are any existing tests that use different tolerances for nopbc cases

# Search for tolerance values in test files
echo "Searching for existing tolerance patterns in tests..."
rg -A 2 'assert_allclose.*nopbc' source/tests/

#!/bin/bash
# Let's search for all assert_allclose usage in test files to understand tolerance patterns
echo "Searching for assert_allclose patterns in tests..."
rg -B 2 -A 2 'assert_allclose' source/tests/consistent/io/test_io.py

# Also search for nopbc related tests
echo -e "\nSearching for nopbc related tests..."
rg -B 2 -A 2 'nopbc' source/tests/consistent/io/test_io.py

#!/bin/bash
# Let's search for tolerance patterns in other test files to get a broader view
echo "Searching for tolerance patterns in all test files..."
rg -B 2 -A 2 'rtol.*atol' source/tests/

# Also search specifically for non-periodic boundary condition related tests
echo -e "\nSearching for other nopbc/non-periodic tests..."
rg -B 2 -A 2 'non.?periodic|nopbc' source/tests/

source/tests/jax/jax2tf/test_nlist.py (2)

34-54: Consider making test data more dynamic and maintainable.

The current setup has several potential improvements:

1. The cell matrix and reference neighbor list are hardcoded, making tests brittle
2. The test could benefit from parameterization to cover more cases

Consider:

1. Using @pytest.mark.parametrize to test different cell configurations
2. Moving test data to external fixtures
3. Adding helper methods to generate reference neighbor lists dynamically

Example parameterization:

@pytest.mark.parametrize("cell_type", [
    "cubic",
    "orthorhombic",
    "triclinic"
])
def test_build_notype(self, cell_type):
    cell = get_test_cell(cell_type)
    # ... rest of the test

---

77-98: Expand type-specific neighbor list testing.

The current test could be enhanced with:

1. More diverse type combinations
2. Boundary cases between different types
3. Verification of type-specific cutoffs

Consider adding:

def test_build_type_diverse(self):
    # Test with more atom types
    diverse_types = tnp.array([-1, 0, 1, 2, 3], dtype=tnp.int32)
    nsel_diverse = [5] * 4  # neighbor selection for each type
    # ... rest of the test

Configuration used: CodeRabbit UI
Review profile: CHILL

source/jax2tf_tests/test_region.py (2)

29-33: Consider documenting the significance of tensor dimensions.

The cell tensor is initialized with specific dimensions [4, 5, -1, 3] but the significance of these dimensions (4 and 5) is not documented. Consider adding comments explaining what these dimensions represent or extracting them as named constants.

---

47-64: Consider enhancing the mathematical documentation.

While the code includes basic comments, it would be beneficial to add more detailed documentation about:

1. The mathematical formulas used for face distance calculations
2. Why volume divided by area gives the correct distance
3. References to relevant geometric principles

This would make the test more maintainable and educational for future contributors.

Example documentation format:

def test_to_face_dist(self):
    """Tests face distance calculations using volume/area method.
    
    Mathematical background:
    - For a parallelepiped, the distance to a face is given by V/A
      where V is the volume and A is the area of the parallel face
    - Volume is calculated using determinant of cell matrix
    - Surface areas are calculated using cross products of edge vectors
    
    References:
    - [Add relevant geometric principles or papers]
    """

source/jax2tf_tests/test_nlist.py (2)

29-55: Add docstrings and comments to explain test data structure.

The setUp method initializes several parameters but lacks documentation explaining:

• The significance of dimensions (nf=3, nloc=3, ns=75)
• The structure of reference neighbor lists
• The coordinate transformation process

Add docstrings following this structure:

def setUp(self):
    """Initialize test data for neighbor list construction.
    
    Test data structure:
    - nf (int): Number of frames
    - nloc (int): Number of local atoms
    - ns (int): Number of shifts (5x5x3 grid)
    - cell (array): Simulation cell matrix
    - icoord (array): Initial coordinates
    - ref_nlist (array): Reference neighbor lists
    """

---

56-76: Enhance test_build_notype with edge cases.

While the basic neighbor list construction is tested, consider adding tests for:

• Empty neighbor lists
• Maximum number of neighbors
• Single atom systems

Configuration used: CodeRabbit UI
Review profile: CHILL